/**
 * 
 */
package solution.tree;

import junit.framework.Assert;
import lib.tree.BNode;
import lib.tree.TreeUtil;

import org.junit.BeforeClass;
import org.junit.Test;

/**
 * @author <a href="www.sureinterview.com">SureInterview</a> Question can be
 *         found at: http://yes.sureinterview.me/shwqst/548002
 */
public class NthSmallestInBST {
	static TreeUtil treeUtil;

	@BeforeClass
	public static void before() {
		treeUtil = new TreeUtil();
	}

	int nodeNumber;

	int nth;

	BNode nthSmallest;

	/**
	 * Find the n-th smallest when the size is unknown.
	 * 
	 * inorder traverse the tree. During the process, bookmark the current node
	 * and return this node when the proper node is found.
	 * 
	 * @param root
	 * @param n
	 * @return
	 */
	BNode find(BNode root, int n) {
		this.nth = n;
		nthSmallest = null;
		nodeNumber = 0;
		findRec(root);
		return nthSmallest;
	}

	/**
	 * Find the n-th smallest when the size of the node is known.
	 * 
	 * 
	 * 
	 * @param root
	 * @param n
	 * @return
	 */
	BNode find2(BNode root, int n) {
		if (root == null || n <= 0 || n > root.size)
			return null;

		int num = 0;
		BNode curNode = root;
		while (num < n && curNode != null) {
			// loop invariant: curNode != null
			// num < n, num + (current node size) > n

			int leftSize = 0;
			if (curNode.left != null) {
				leftSize = curNode.left.size;
			}

			if (leftSize + num >= n) {
				// the node falls into left tree
				curNode = curNode.left;
				continue;
			} else if (leftSize + 1 + num == n)
				// current node is what we want.
				return curNode;
			else {
				// the node falls into right tree.
				curNode = curNode.right;
				num += leftSize + 1;
			}
		}
		return null;
	}

	void findRec(BNode node) {
		// 1. setup/precondition
		if (node == null || nthSmallest != null)
			return;

		// 2.1 left
		findRec(node.left);
		if (nthSmallest != null)
			return;

		/**********************
		 * check the current position, return this node when it is the n-th
		 * smallest.
		 */
		// 2.2 root
		nodeNumber++;
		if (nodeNumber == nth) {
			nthSmallest = node;
			return;
		}
		/**********************/

		// 2.3 right
		findRec(node.right);

		// 3 clean up
	}

	@Test
	public void test() {
		BNode tree = treeUtil.makeTree("4(2(1,3),6(5,7)))");
		for (Integer i = 1; i < 7; i++) {
			Assert.assertEquals(i, find(tree, i).value);
			Assert.assertEquals(i, find2(tree, i).value);
		}
	}
}
